using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
import ModelingToolkit: Interval, infimum, supremum

@parameters x t
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2

# Thermal conductivity constant
k = 50 # replace with actual value

# Gaussian beam source term for 1D
P = 10 # peak power of laser, replace with actual value
ω = 0.001 # beam waist, adjust for 1D

# Physical properties of steel
lambda_pr = 0.8
rho_steel = 7850  # Density of steel in kg/m³
cp_steel = 500    # Specific heat capacity in J/(kg·K)
k_steel = 50      # Thermal conductivity in W/(m·K)

gaussian_beam_1d(x) = ((2 * lambda_pr * P)/(pi*ω^2)) * exp(-2 * (x^2 / ω^2)) #
#gaussian_beam_1d(x) = ((2 * lambda_pr * P)/(pi*ω^2)) * exp(-2 * ((x - x0)^2 / ω^2))

# Heat equation for 1D
eq = Dt(u(x, t)) ~ k * Dxx(u(x, t)) - gaussian_beam_1d(x)

# Boundary and initial conditions for 1D
#bcs = [u(0, t) ~ 0.0, u(1, t) ~ 0, u(x, 0) ~ 0]
# Boundary and initial conditions for 1D with surrounding and initial temperature at 20 Celsius
ambient_temp = 20.0 # Ambient temperature in Celsius

bcs = [u(0, t) ~ ambient_temp, u(1, t) ~ ambient_temp, u(x, 0) ~ ambient_temp]

# The rest of your code remains unchanged...

# Space and time domains
domains = [x ∈ Interval(0.0, 1.0), t ∈ Interval(0.0, 1.0)]

# Discretization parameters for 1D
dx = 0.1
dt = 0.1

# Neural network architecture for 1D
chain = Lux.Chain(Dense(2, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))

# Physics-Informed Neural Network for 1D
discretization = PhysicsInformedNN(chain, QuadratureTraining())

@named pde_system = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])
prob = discretize(pde_system, discretization)

# Callback for loss in 1D
callback = function (p, l)
    println("Current loss is: $l")
    return false
end

# Existing setup for the Physics-Informed Neural Network (PINN)
# ...

# Define the analytical solution function
function analytic_sol_func(t)
    return 20.0 + (2 * F0 / k_steel) * (k_th_diff * t / pi)^0.5
end

# Modify the loss function to include the analytical solution
function custom_loss_function(phi, p, t, analytic_sol_func)
    # Prediction from the neural network
    nn_prediction = phi(p)

    # Evaluation of the analytical solution
    analytic_prediction = analytic_sol_func(t)

    # Calculate the mean squared error between NN prediction and analytical solution
    analytical_loss = mean((nn_prediction .- analytic_prediction).^2)

    # Combine this loss with your existing loss function
    total_loss = existing_loss + analytical_loss

    return total_loss
end

# Solving the PDE in 1D
res = Optimization.solve(prob, ADAM(0.1); callback = callback, loss_function = custom_loss_function, maxiters = 4000)
prob = remake(prob, u0 = res.minimizer)
res = Optimization.solve(prob, ADAM(0.01); callback = callback, maxiters = 2000)
phi = discretization.phi

# Assuming the rest of your code is unchanged
# Prediction and plotting for 1D
xs, ts = [infimum(d.domain):(dx / 10):supremum(d.domain) for d in domains]
u_predict_1d = reshape([first(phi([x, t], res.minimizer)) for x in xs for t in ts],
                       (length(xs), length(ts)))

# Find the index for x = 0.1
x_val = 0.0
x_index = findfirst(≈(x_val), xs)

# Extract temperature values at x = 0.1 over time
temp_at_x = u_predict_1d[x_index, :]

# Plot temperature over time at x = 0.1
using Plots

using JSON

#thermal diffusivity:
k_th_diff = k_steel/(rho_steel*cp_steel)
F0 = P/pi*ω^2

analytic_sol_func(t) = 20.0 + (2*F0/k_steel)*(k_th_diff*t/pi)^0.5

u_real = reshape([analytic_sol_func(t) for t in ts], length(ts))

# Assuming temp_at_x and u_real are arrays containing your data
data = Dict("temp_at_x" => temp_at_x, "u_real" => u_real, "ts" => ts,)

json_data = JSON.json(data)

open("data.json", "w") do f
    write(f, json_data)
end

p1=plot(ts, temp_at_x, xlabel = "Time", ylabel = "Temperature", title = "Temperature at x = $x_val Over Time", legend = false)



# Plot the analytical solution over time
p2=plot(ts, u_real, xlabel = "Time", ylabel = "Temperature", title = "Analytical Solution Over Time", legend = false)
plot(p1, p2)

